今天我們來了解如何使用pl.DataFrame.group_by()，進行聚合運算。
本日大綱如下：
pl.Expr.over()
codepanda
import pandas as pd
import polars as pl
data = {
"name": ["Tom", "Lisa", "John", "Vincent", "Mary", "Caroline"],
"has_pet": ["Y", "N", "Y", "Y", "Y", "N"],
"gender": ["M", "F", "M", "M", "F", "F"],
"lucky_number": [19, 25, 36, 7, 2, 91],
}
df = pl.DataFrame(data)
shape: (6, 4)
┌──────────┬─────────┬────────┬──────────────┐
│ name ┆ has_pet ┆ gender ┆ lucky_number │
│ --- ┆ --- ┆ --- ┆ --- │
│ str ┆ str ┆ str ┆ i64 │
╞══════════╪═════════╪════════╪══════════════╡
│ Tom ┆ Y ┆ M ┆ 19 │
│ Lisa ┆ N ┆ F ┆ 25 │
│ John ┆ Y ┆ M ┆ 36 │
│ Vincent ┆ Y ┆ M ┆ 7 │
│ Mary ┆ Y ┆ F ┆ 2 │
│ Caroline ┆ N ┆ F ┆ 91 │
└──────────┴─────────┴────────┴──────────────┘
pl.DataFrame.group_by().agg()為聚合的基本型式：
group_by()內可以置入一或多個列名/expr做為分組依據。agg()內為一或多個expr，代表聚合的各種操作。舉例來說，我們可以針對「"has_pet"」列進行分組，並使用pl.len()計算各組所包含的人數及找出各組中「"lucky_number"」列中最大的數值（註1）：
df.group_by("has_pet").agg(pl.len(), pl.max("lucky_number"))
shape: (2, 3)
┌─────────┬─────┬──────────────┐
│ has_pet ┆ len ┆ lucky_number │
│ --- ┆ --- ┆ --- │
│ str ┆ u32 ┆ i64 │
╞═════════╪═════╪══════════════╡
│ Y ┆ 4 ┆ 36 │
│ N ┆ 2 ┆ 91 │
└─────────┴─────┴──────────────┘
此外，使用expr做為分組依據也是很常見的應用。例如我們想要以「"name"」列中各個名字的長度來做為分組依據，並計算各組所包含的人數以及收集各組所包含的名字，可以這麼寫：
(
df.group_by(pl.col("name").str.len_bytes().alias("n_chars"))
.agg(pl.len(), pl.col("name"))
.sort("n_chars")
)
shape: (4, 3)
┌─────────┬─────┬──────────────────────────┐
│ n_chars ┆ len ┆ name │
│ --- ┆ --- ┆ --- │
│ u32 ┆ u32 ┆ list[str] │
╞═════════╪═════╪══════════════════════════╡
│ 3 ┆ 1 ┆ ["Tom"] │
│ 4 ┆ 3 ┆ ["Lisa", "John", "Mary"] │
│ 7 ┆ 1 ┆ ["Vincent"] │
│ 8 ┆ 1 ┆ ["Caroline"] │
└─────────┴─────┴──────────────────────────┘
此處使用了pl.col("name").str.len_bytes().alias("n_chars")這個expr來計算名字長度。其中agg()中的pl.col("name")（也可以直接使用列名「"name"」）可以將該組的結果收集為一個pl.List。請注意，為了維持呈現結果，最後使用了pl.DataFrame.sort()將結果以名字的長度由短至長排序。
對於分組後的聚合操作，我們可能會想要依據某些條件，將結果呈現於不同列。例如想針對「"gender"」列進行分組，並將結果分為兩列，一列是各組中「"lucky_number"」小於20的人數，而另一列則是各組中「"lucky_number"」大於或等於20的人數（註2）：
lt20 = pl.col("lucky_number").lt(20)
(
df.group_by("gender").agg(
lt20.sum().alias("lt20"), lt20.not_().sum().alias("~lt20")
)
)
shape: (2, 3)
┌────────┬──────┬───────┐
│ gender ┆ lt20 ┆ ~lt20 │
│ --- ┆ --- ┆ --- │
│ str ┆ u32 ┆ u32 │
╞════════╪══════╪═══════╡
│ M ┆ 2 ┆ 1 │
│ F ┆ 1 ┆ 2 │
└────────┴──────┴───────┘
另一種使用方式是使用pl.Expr.filter()針對各分組結果再進行篩選。例如針對「"gender"」列進行分組，並想要收集各分組內「"lucky_number"」小於20的名字，可以這麼寫：
(
df.group_by("gender").agg(
pl.col("name").filter(pl.col("lucky_number").lt(20))
)
)
shape: (2, 2)
┌────────┬────────────────────┐
│ gender ┆ name │
│ --- ┆ --- │
│ str ┆ list[str] │
╞════════╪════════════════════╡
│ M ┆ ["Tom", "Vincent"] │
│ F ┆ ["Mary"] │
└────────┴────────────────────┘
我們也可以針對多列或是多個expr進行分組聚合。例如此處我們針對「"has_pet"」列及「"gender"」列進行分組，並收集各組名字為一Pl.List：
(df.group_by("has_pet", "gender", maintain_order=True).agg(pl.col("name")))
shape: (3, 3)
┌─────────┬────────┬────────────────────────────┐
│ has_pet ┆ gender ┆ name │
│ --- ┆ --- ┆ --- │
│ str ┆ str ┆ list[str] │
╞═════════╪════════╪════════════════════════════╡
│ Y ┆ M ┆ ["Tom", "John", "Vincent"] │
│ N ┆ F ┆ ["Lisa", "Caroline"] │
│ Y ┆ F ┆ ["Mary"] │
└─────────┴────────┴────────────────────────────┘
為了方便說明，我們將maintain_order=設為True，並逐一講解如何將原先的六列資料組合為三組：
針對一些常用的聚合，Polars提供了一個便捷的寫法，詳情可以參考API文件。假如我們想要針對「"has_pet"」列進行分組，並計算各組所包含的人數，除了我們已經熟悉的agg()搭配pl.len()外：
df.group_by("has_pet", maintain_order=True).agg(pl.len())
也可以這麼寫：
df.group_by("has_pet", maintain_order=True).len()
兩者會得到一樣的結果：
shape: (2, 2)
┌─────────┬─────┐
│ has_pet ┆ len │
│ --- ┆ --- │
│ str ┆ u32 │
╞═════════╪═════╡
│ Y ┆ 4 │
│ N ┆ 2 │
└─────────┴─────┘
pl.Expr.over()pl.Expr.over()中提到，這個expr的功能就像是PostgreSQL中的window function。
簡單地說，pl.Expr.over()是可以針對分組後資料進行個別運算的expr。舉例來說，假如我們想針對「"gender"」列進行分組運算，計算每組內「"lucky_number"」列的rank（即排序各組內的「"lucky_number"」列，最小值為1，次小值為2，以此類推），可以這麼寫（註3）：
(
df.with_columns(
pl.col("lucky_number")
.rank("ordinal")
.over("gender")
.alias("rank_by_gender")
)
)
shape: (6, 5)
┌──────────┬─────────┬────────┬──────────────┬────────────────┐
│ name ┆ has_pet ┆ gender ┆ lucky_number ┆ rank_by_gender │
│ --- ┆ --- ┆ --- ┆ --- ┆ --- │
│ str ┆ str ┆ str ┆ i64 ┆ u32 │
╞══════════╪═════════╪════════╪══════════════╪════════════════╡
│ Tom ┆ Y ┆ M ┆ 19 ┆ 2 │
│ Lisa ┆ N ┆ F ┆ 25 ┆ 2 │
│ John ┆ Y ┆ M ┆ 36 ┆ 3 │
│ Vincent ┆ Y ┆ M ┆ 7 ┆ 1 │
│ Mary ┆ Y ┆ F ┆ 2 ┆ 1 │
│ Caroline ┆ N ┆ F ┆ 91 ┆ 3 │
└──────────┴─────────┴────────┴──────────────┴────────────────┘
可以看出，「"rank_by_gender"」列總共有兩組「123」，分別代表該行在各自性別的排序。
以第一行為例，根據「"gender"」列來看，「"Tom"」屬於「"M"」群組，其「"lucky_number"」為「19」，介於同組的「"Vincent"」及「"John"」之間，所以其「"rank_by_gender"」為2。
事實上，如果使用pl.DataFrame.group_by().agg()搭配pl.DataFrame.explode()，可以做出類似上面的結果：
(
df.group_by("gender", maintain_order=True)
.agg(
pl.col("lucky_number").rank("ordinal").alias("rank_by_gender"),
pl.all(),
)
.explode(pl.all().exclude("gender"))
.select([*df.columns, "rank_by_gender"])
)
shape: (6, 5)
┌──────────┬─────────┬────────┬──────────────┬────────────────┐
│ name ┆ has_pet ┆ gender ┆ lucky_number ┆ rank_by_gender │
│ --- ┆ --- ┆ --- ┆ --- ┆ --- │
│ str ┆ str ┆ str ┆ i64 ┆ u32 │
╞══════════╪═════════╪════════╪══════════════╪════════════════╡
│ Tom ┆ Y ┆ M ┆ 19 ┆ 2 │
│ John ┆ Y ┆ M ┆ 36 ┆ 3 │
│ Vincent ┆ Y ┆ M ┆ 7 ┆ 1 │
│ Lisa ┆ N ┆ F ┆ 25 ┆ 2 │
│ Mary ┆ Y ┆ F ┆ 2 ┆ 1 │
│ Caroline ┆ N ┆ F ┆ 91 ┆ 3 │
└──────────┴─────────┴────────┴──────────────┴────────────────┘
可以觀察出來，其「"gender"」會是同組一起呈現，而不是像前一個例題那樣交叉呈現。
此外，pl.Expr.over()有一個重要的mapping_strategy=參數，總共有「"group_to_rows"」、「"explode"」及 「"join"」三種選項（預設為「"group_to_rows"」），可以改變呈現結果。以下我們建立一個df2 dataframe幫助說明，其內除了有df的「"name"」及「"gender"」列，還新增了一列「"rank"」列。請注意「"rank"」列內共有六個數字，代表這是全局的排列順序，不是如前面例題的個別分組排序。
df2 = df.with_columns(
pl.Series("rank", [5, 6, 4, 1, 2, 3], dtype=pl.UInt32)
).select("name", "gender", "rank")
shape: (6, 3)
┌──────────┬────────┬──────┐
│ name ┆ gender ┆ rank │
│ --- ┆ --- ┆ --- │
│ str ┆ str ┆ u32 │
╞══════════╪════════╪══════╡
│ Tom ┆ M ┆ 5 │
│ Lisa ┆ F ┆ 6 │
│ John ┆ M ┆ 4 │
│ Vincent ┆ M ┆ 1 │
│ Mary ┆ F ┆ 2 │
│ Caroline ┆ F ┆ 3 │
└──────────┴────────┴──────┘
以下我們使用pl.Expr.over()針對「"gender"」列分組，並依各組內的「"rank"」排序後，觀察三種mapping_strategy=參數的結果有什麼不同。
「"group_to_rows"」將分組後的結果「置回」原先的dataframe，這也是其參數名的由來，即將「"group"」置入「"rows"」中。
(
df2.select(
pl.all()
.sort_by(pl.col("rank"))
.over(pl.col("gender"), mapping_strategy="group_to_rows"),
)
)
shape: (6, 3)
┌──────────┬────────┬──────┐
│ name ┆ gender ┆ rank │
│ --- ┆ --- ┆ --- │
│ str ┆ str ┆ u32 │
╞══════════╪════════╪══════╡
│ Vincent ┆ M ┆ 1 │
│ Mary ┆ F ┆ 2 │
│ John ┆ M ┆ 4 │
│ Tom ┆ M ┆ 5 │
│ Caroline ┆ F ┆ 3 │
│ Lisa ┆ F ┆ 6 │
└──────────┴────────┴──────┘
「"explode"」的效果與「"group_to_rows"」很像，但是其結果會是同組一起呈現。教學文件特別指出，如果各組內之行排序不重要時，應優先使用「"explode"」，因為這樣Polars可以減少「"group to rows"」的操作。
(
df2.select(
pl.all()
.sort_by(pl.col("rank"))
.over(pl.col("gender"), mapping_strategy="explode")
)
)
shape: (6, 3)
┌──────────┬────────┬──────┐
│ name ┆ gender ┆ rank │
│ --- ┆ --- ┆ --- │
│ str ┆ str ┆ u32 │
╞══════════╪════════╪══════╡
│ Vincent ┆ M ┆ 1 │
│ John ┆ M ┆ 4 │
│ Tom ┆ M ┆ 5 │
│ Mary ┆ F ┆ 2 │
│ Caroline ┆ F ┆ 3 │
│ Lisa ┆ F ┆ 6 │
└──────────┴────────┴──────┘
「"join"」則會將各組結果合成一個pl.List，並重覆出現在該組。API文件特別指出，這是一個耗費記憶體的操作，需小心使用。
(
df2.with_columns(
pl.col("rank")
.sort()
.over(pl.col("gender"), mapping_strategy="join")
)
)
shape: (6, 3)
┌──────────┬────────┬───────────┐
│ name ┆ gender ┆ rank │
│ --- ┆ --- ┆ --- │
│ str ┆ str ┆ list[u32] │
╞══════════╪════════╪═══════════╡
│ Tom ┆ M ┆ [1, 4, 5] │
│ Lisa ┆ F ┆ [2, 3, 6] │
│ John ┆ M ┆ [1, 4, 5] │
│ Vincent ┆ M ┆ [1, 4, 5] │
│ Mary ┆ F ┆ [2, 3, 6] │
│ Caroline ┆ F ┆ [2, 3, 6] │
└──────────┴────────┴───────────┘
最後特別提醒，雖然pl.Expr.over()大部份情況會產生跟原先dataframe一樣的行數，但這不是一定的。舉例來說，我們可以透過pl.Expr.head()只取得各「"gender"」組內最小「"rank"」的行。也就是說在「"M"」群組中，取到該組「"rank"」最小的「"Vincent"」，而在「"F"」群組中，取到該組「"rank"」最小的「"Mary"」。
(
df2.select(
pl.all()
.sort_by(pl.col("rank"))
# for each gender, get the first row after sorting by "rank"
.head(1)
.over(pl.col("gender"), mapping_strategy="explode")
)
)
shape: (2, 3)
┌─────────┬────────┬──────┐
│ name ┆ gender ┆ rank │
│ --- ┆ --- ┆ --- │
│ str ┆ str ┆ u32 │
╞═════════╪════════╪══════╡
│ Vincent ┆ M ┆ 1 │
│ Mary ┆ F ┆ 2 │
└─────────┴────────┴──────┘
請注意，此處的pl.Expr.head(1)是在pl.Expr.over()之前，代表其是針對每一組而操作。
如果是將pl.Expr.head(1)置於pl.Expr.over()後，代表進行完pl.Expr.over()操作後，取得第一行：
(
df2.select(
pl.all()
.sort_by(pl.col("rank"))
.over(pl.col("gender"), mapping_strategy="explode")
.head(1)
)
)
shape: (1, 3)
┌─────────┬────────┬──────┐
│ name ┆ gender ┆ rank │
│ --- ┆ --- ┆ --- │
│ str ┆ str ┆ u32 │
╞═════════╪════════╪══════╡
│ Vincent ┆ M ┆ 1 │
└─────────┴────────┴──────┘
codepandaPandas中相對應於polars的pl.DataFrame.group_by()的功能是pd.DataFrame.groupby()。
例如想要針對「"has_pet"」列進行分組，並計算各組所包含的人數及找出各組中「"lucky_number"」列中最大的數值，可以這麼寫：
df_pd = pd.DataFrame(data)
(
df_pd.groupby("has_pet").agg(
len=("has_pet", "size"),
lucky_number=("lucky_number", "max"),
)
)
len lucky_number
has_pet
N 2 91
Y 4 36
如果想要以「"name"」列中各個名字的長度來做為分組依據，並計算各組所包含的人數以及收集各組所包含的名字，可以這麼寫：
(
df_pd.assign(n_chars=df_pd["name"].str.len())
.groupby("n_chars")
.agg(len=("name", "size"), name=("name", list))
)
len name
n_chars
3 1 [Tom]
4 3 [Lisa, John, Mary]
7 1 [Vincent]
8 1 [Caroline]
Pandas中最接近polarspl.Expr.over()的功能是pd.DataFrame.groupby.DataFrameGroupBy.transform()。
舉例來說，假如我們想針對「"gender"」列進行分組運算，計算每組內「"lucky_number"」列的rank（即排序各組內的「"lucky_number"」列，最小值為1，次小值為2，以此類推），可以這麼寫：
(
df_pd.assign(
rank_by_gender=lambda df_: df_.groupby("gender")
.lucky_number.transform(lambda s_: s_.rank())
.astype(int)
)
)
name has_pet gender lucky_number rank_by_gender
0 Tom Y M 19 2
1 Lisa N F 25 2
2 John Y M 36 3
3 Vincent Y M 7 1
4 Mary Y F 2 1
5 Caroline N F 91 3
註1：多次執行相同的分組聚合運算時，其各行順序不一定會一樣，也就是可能會出現「"N"」出現在「"Y"」前面的情況，例如：
shape: (2, 3)
┌─────────┬─────┬──────────────┐
│ has_pet ┆ len ┆ lucky_number │
│ --- ┆ --- ┆ --- │
│ str ┆ u32 ┆ i64 │
╞═════════╪═════╪══════════════╡
│ N ┆ 2 ┆ 91 │
│ Y ┆ 4 ┆ 36 │
└─────────┴─────┴──────────────┘
想要維持各行為固定順序的話，有兩個方法：
pl.DataFrame.group_by()的maintain_order=設為True，這樣一來Polars就會以DataFrame中的順序為依據來呈現結果。註2：這裡使用了pl.Expr.not_()來做反向（negate）布林運算。如果不習慣這樣寫法的朋友，可以自行拼湊出反向的expr，如：
lt20 = pl.col("lucky_number").lt(20)
ge20 = pl.col("lucky_number").ge(20)
(
df.group_by("gender").agg(
lt20.sum().alias("lt20"), ge20.sum().alias("ge20")
)
)
shape: (2, 3)
┌────────┬──────┬──────┐
│ gender ┆ lt20 ┆ ge20 │
│ --- ┆ --- ┆ --- │
│ str ┆ u32 ┆ u32 │
╞════════╪══════╪══════╡
│ M ┆ 2 ┆ 1 │
│ F ┆ 1 ┆ 2 │
└────────┴──────┴──────┘
註3：pl.Expr.over()的第一個參數method=，是為了判斷當兩個元素相等時，如何決定大小所用，其預設值為「"average"」，型別會是Pl.Float64。由於我們的例題中不會出現這樣的情況，所以我選擇指定method=為ordinal，型別為pl.UInt32。
iThome鐵人賽
看影片追技術
看更多